# -*- coding: utf-8 -*-
from naie.datasets import get_data_reference
from naie.context import Context
import moxing as mox

import zipfile
import os
import time

import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
from torchvision import models

# from tensorboardX import SummaryWriter
import pandas as pd
from PIL import Image
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True


class MyDataset(Dataset):
    def __init__(self, root, input_data, aug):
        self.file_data = input_data['FileID'].values
        self.label_data = input_data[
            'SpeciesID'].values if 'SpeciesID' in input_data.columns else None
        self.aug = aug
        self.img_data = [Image.open(root + i + '.jpg') for i in self.file_data]

    def __len__(self):
        return len(self.img_data)

    def __getitem__(self, index):
        img = self.img_data[index]

        if self.aug is not None:
            img = self.aug(img)

        if self.label_data is not None:
            return img, self.file_data[index], self.label_data[index]
        else:
            return img, self.file_data[index]


def test():
    model.eval()
    ans_file = []
    ans_pred = []
    for _, (img, fileid) in enumerate(test_dataloader):
        inputs = img.to(device).float()
        outputs = model(inputs)
        ans_file.extend(fileid)
        ans_pred.extend(outputs.max(1)[1].detach().cpu().numpy())
    ans = [[ans_file[i], ans_pred[i]] for i in range(len(ans_file))]
    ans = pd.DataFrame(ans, columns=['FileID', 'SpeciesID'])
    ans.to_csv('ans.csv', index=None)
    mox.file.copy('ans.csv', os.path.join(Context.get_output_path(), 'ans.csv'))
    print('ans saved.')


def load_model_state_dict_and_online_dataset():
    # load model state dict
    data_reference = get_data_reference(dataset="pretrain",
                                        dataset_entity='resnet50_best25')
    for file_path in data_reference.get_files_paths():
        mox.file.copy(file_path, './cache/' + file_path.split('/')[-1])

    # load data
    data_reference = get_data_reference(dataset="DatasetService",
                                        dataset_entity='image_dataset')
    for file_path in data_reference.get_files_paths():
        mox.file.copy(file_path, './cache/' + file_path.split('/')[-1])
    zip_file = zipfile.ZipFile('./cache/data.zip')
    zip_list = zip_file.namelist()
    for f in zip_list:
        zip_file.extract(f, './cache/')


if __name__ == "__main__":    

    # 载入数据集和模型参数
    load_model_state_dict_and_online_dataset()
    
    # 设置图片文件夹路径，读取相应csv文件，设置欲载入的模型断点文件路径
    img_path = './cache/data/'
    test_csv = pd.read_csv('./cache/test.csv')
    mox.file.copy('annotation.csv', './cache/annotation.csv')
    valid_csv = pd.read_csv('./cache/annotation.csv')
    LOAD_CHECKPOINT_PATH = r'./cache/resnet50-25-best.pth'

    # 硬件选择，批大小，线程数
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    BATCH_SIZE, NUM_WORKERS = 32, 4
    print('device:{} batch_size:{} num_workers:{}'.format(device, BATCH_SIZE, NUM_WORKERS))

    # 测试图片数据预处理
    test_aug = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225]),
    ])

    # 测试数据导入
    test_data = MyDataset(root=img_path, input_data=test_csv, aug=test_aug)
    test_dataloader = DataLoader(dataset=test_data,
                                 batch_size=BATCH_SIZE,
                                 shuffle=False,
                                 num_workers=NUM_WORKERS)

    print('num_test_data:{}'.format(len(test_data)))

    # 模型
    # from cache.resnet50 import resnet50  # 严格按照路径进行import！
    model = models.resnet50(pretrained=True)
    fc_features = model.fc.in_features
    model.fc = nn.Linear(fc_features, 20)
    model = model.to(device)
    modelname = str(model.__class__).strip('<>').split()[-1].strip("'").split('.')[-1]
    print(modelname)

    # 载入模型断点参数
    if torch.cuda.is_available():
        loaded_checkpoint = torch.load(LOAD_CHECKPOINT_PATH)
    else:
        loaded_checkpoint = torch.load(LOAD_CHECKPOINT_PATH, map_location=torch.device('cpu'))
    model.load_state_dict(loaded_checkpoint)
    print('load checkpoint: {}'.format(LOAD_CHECKPOINT_PATH))

    # 开始测试
    print('Start testing...')
    with torch.no_grad():
        test()
    print('Testing Task Finished.')

os.system('python evaluate_ans.py')
